import PIL
import numpy as np
from PIL import Image
import os
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor,Compose
import torch

mnist_dataset_train=MNIST(root="./data",train=True,download=True,
                          transform= Compose([ToTensor()]) )
mnist_dataset_test=MNIST(root="./data",train=False ,download=True,
                         transform= Compose([ToTensor()]))

# shape
# for i, data in enumerate(mnist_dataset_train):
#     print(f"{i},{type(data)}") #(img,label)
#     # print(data)
#     print(data[0].shape,data[1]) #img [1,28,28]
#     break

# show
# data=mnist_dataset_train[0]
# img=data[0].numpy() #1
# print(img.shape)
# img=Image.fromarray(img.squeeze()*255) #2
# print(img.size)
# img.show()

# save
data_path="./data/MNIST/"
train_path=data_path+"/train"
test_path=data_path+"/test"

with open(data_path+'train.txt','w') as f:
    for i, data in enumerate(mnist_dataset_train):
        img_name=str(i)
        img_name=img_name.zfill(6)
        img_path=os.path.join(train_path,img_name+".png")
        print(img_path," ", data[1])
        #
        img=data[0].numpy() #1
        img=img*255 #2
        img=img.astype(np.uint8) #3
        img=Image.fromarray(img.squeeze()) #4
        img.save(img_path)
        # Image.fromarray((data[0].numpy().squeeze() * 255).astype(np.uint8)).save(img_path)
        f.write(img_path+" "+str(data[1])+"\n")

with open(data_path+'test.txt','w') as f:
    for i, data in enumerate(mnist_dataset_test):
        img_name=str(i)
        img_name=img_name.zfill(6)
        img_path=os.path.join(test_path,img_name+".png")
        print(img_path," ", data[1])
        #
        img=data[0].numpy() #1
        img=img*255 #2
        img=img.astype(np.uint8) #3
        img=Image.fromarray(img.squeeze()) #4
        img.save(img_path)
        f.write(img_path+" "+str(data[1])+"\n")

# pass
with open(data_path+'train.txt','r') as f:
    lines=f.readlines()
for line in lines:
    data=line.split(" ")
    # print(data)
    img_path, label =data[0],data[1]
    print(img_path, label)

img=Image.open(img_path)
img.show()